//
//  RedBlackTree.hpp
//  封装map和set
//
//  Created by 卜绎皓 on 2022/11/17.
//

#pragma once
#include<iostream>
using namespace std;
 
enum Colour
{
    RED,
    BLACK
};//节点的颜色
template<class T>
struct RBTreeNode
{
    RBTreeNode<T>* _left;
    RBTreeNode<T>* _right;
    RBTreeNode<T>* _parent;
    T _data;
    Colour _col;
    RBTreeNode(const T& data)
        :_left(nullptr)
        , _right(nullptr)
        , _parent(nullptr)
        , _data(data)
        , _col(RED)
    {}
};
template<class T, class Ref, class Ptr>
struct __TreeIterator//封装迭代器
{
    typedef RBTreeNode<T> Node;//迭代器就是节点的指针
    typedef __TreeIterator<T, Ref, Ptr> Self;
 
    Node* _node;//
    __TreeIterator(Node* node)
        :_node(node)
    {}
    Ref operator*()
    {
        return _node->_data;
    }
    Ptr operator->()
    {
        return &_node->_data;
    }
    bool operator != (const Self& s) const
    {
        return _node != s._node;
    }
    bool operator == (const Self& s) const
    {
        return _node == s._node;
    }
    // 迭代器的难点
    Self& operator++()
    {
        //++注意画图分析
        if (_node->_right)//节点的右子树不为空
        {
            // 下一个访问就是右树中，中序的第一个节点
            Node* left = _node->_right;
            while (left->_left)//循环去找最左节点
            {
                left = left->_left;
            }
            _node = left;
        }
        else
        {
            // 找祖先里面不是父亲的右的那个节点
            // 因为 cur 右为空，说明cur所在的子树已经访问完了
            // cur是parent的右的，说明parent包括其所在子树都访问完了，循环往上去找
            Node* cur = _node;
            Node* parent = cur->_parent;
            while (parent && cur == parent->_right)//如果父亲不为空或者cur为父亲的右节点
            {
                cur = cur->_parent;
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }
    Self& operator--()
    {
        if (_node->_left)
        {
            // 下一个需要访问左子树的最右节点
            Node* right = _node->_left;
            while (right->_right)
            {
                right = right->_right;
            }
            _node = right;
        }
        else
        {
            Node* cur = _node;
            Node* parent = cur->_parent;
            //如果父节点存在且cur一直为父节点的左节点循环继续向上找
            while (parent && cur == parent->_left)
            {
                cur = parent;
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }
};
template<class K, class T, class KeyOfT>//多添加一个模板参数用以map和set进行封装
class RBTree
{
    typedef RBTreeNode<T> Node;
public:
    typedef __TreeIterator < T, T&, T* > iterator;
    typedef __TreeIterator < T, const T&, const T* > const_iterator;
 
    iterator begin()
    {
        Node* left = _root;
        while (left && left->_left)
        {
            left = left->_left;
        }
        return iterator(left);
    }
    iterator end()
    {
        return iterator(nullptr);
    }
    RBTree()
        :_root(nullptr)
    {}
    void Destory(Node* root)
    {
        if (root == nullptr)
        {
            return;
        }
 
        Destory(root->_left);
        Destory(root->_right);
        delete root;
    }
    ~RBTree()
    {
        Destory(_root);
        _root = nullptr;
    }
    Node* Find(const K& key)
    {
        KeyOfT kot;
        Node* cur = _root;
        while (cur)
        {
            if (kot(cur->_data) > key)
            {
                cur = cur->_left;
            }
            else if (kot(cur->_data) < key)
            {
                cur = cur->_right;
            }
            else
            {
                return cur;
            }
        }
        return nullptr;
    }
    pair<iterator, bool> Insert(const T& data)
    {
        if (_root == nullptr)
        {
            _root = new Node(data);
            _root->_col = BLACK;
            return make_pair(iterator(_root), true);
        }
        KeyOfT kot;//仿函数对象,重载operator()即可
        Node* parent = nullptr;
        Node* cur = _root;
        while (cur)
        {
            if (kot(cur->_data) < kot(data))
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (kot(cur->_data) > kot(data))
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return make_pair(iterator(cur), false);
            }
        }
        Node* newnode = new Node(data);
        newnode->_col = RED;
        if (kot(parent->_data) < kot(data))
        {
            parent->_right = newnode;
            newnode->_parent = parent;
        }
        else
        {
            parent->_left = newnode;
            newnode->_parent = parent;
        }
        cur = newnode;
 
        // 如果父亲存在，且颜色为红色就需要处理
        while (parent && parent->_col == RED)
        {
            Node* grandfather = parent->_parent;
            // 关键是看叔叔
            if (parent == grandfather->_left)
            {
                Node* uncle = grandfather->_right;
                // 情况1：uncle存在且为红
                if (uncle && uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLACK;
                    grandfather->_col = RED;
 
                    // 继续往上处理
                    cur = grandfather;
                    parent = cur->_parent;
                }
                else // 情况2+3：uncle不存在 uncle存在且为黑
                {
                    // 情况2：单旋
                    if (cur == parent->_left)
                    {
                        RotateR(grandfather);
                        grandfather->_col = RED;
                        parent->_col = BLACK;
                    }
                    else // 情况3：双旋
                    {
                        RotateL(parent);
                        RotateR(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
 
                    break;
                }
            }
            else // parent == grandfather->_right
            {
                Node* uncle = grandfather->_left;
                // 情况1：
                if (uncle && uncle->_col == RED)
                {
                    uncle->_col = parent->_col = BLACK;
                    grandfather->_col = RED;
 
                    cur = grandfather;
                    parent = cur->_parent;
                }
                else // 情况2：+ 情况3：
                {
                    if (cur == parent->_right)
                    {
                        RotateL(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    else // cur == parent->_left
                    {
                        RotateR(parent);
                        RotateL(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    // 插入结束
                    break;
                }
            }
        }
 
        _root->_col = BLACK;
        return make_pair(iterator(newnode), true);
    }
 
    void RotateL(Node* parent)
    {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
 
        parent->_right = subRL;
        if (subRL)
        {
            subRL->_parent = parent;
        }
 
        subR->_left = parent;
        Node* parentParent = parent->_parent;
        parent->_parent = subR;
 
        if (parent == _root)
        {
            _root = subR;
            _root->_parent = nullptr;
        }
        else
        {
            if (parentParent->_left == parent)
            {
                parentParent->_left = subR;
            }
            else
            {
                parentParent->_right = subR;
            }
            subR->_parent = parentParent;
        }
    }
 
    void RotateR(Node* parent)
    {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
 
        parent->_left = subLR;
        if (subLR)
            subLR->_parent = parent;
 
        subL->_right = parent;
        Node* parentParent = parent->_parent;
        parent->_parent = subL;
 
        if (parent == _root)
        {
            _root = subL;
            _root->_parent = nullptr;
        }
        else
        {
            if (parentParent->_left == parent)
                parentParent->_left = subL;
            else
                parentParent->_right = subL;
 
            subL->_parent = parentParent;
        }
    }
 
    bool _CheckBlance(Node* root, int blackNum, int count)
    {
        if (root == nullptr)
        {
            if (count != blackNum)
            {
                cout << "黑色节点的数量不相等" << endl;
                return false;
            }
 
            return true;
        }
 
        if (root->_col == RED && root->_parent->_col == RED)
        {
            cout << "存在连续的红色节点" << endl;
            return false;
        }
 
        if (root->_col == BLACK)
        {
            count++;
        }
 
        return _CheckBlance(root->_left, blackNum, count)
            && _CheckBlance(root->_right, blackNum, count);
    }
 
    bool CheckBlance()
    {
        if (_root == nullptr)
        {
            return true;
        }
 
        if (_root->_col == RED)
        {
            cout << "根节点是红色的" << endl;
            return false;
        }
 
        // 找最左路径做黑色节点数量参考值
        int blackNum = 0;
        Node* left = _root;
        while (left)
        {
            if (left->_col == BLACK)
            {
                blackNum++;
            }
 
            left = left->_left;
        }
 
        int count = 0;
        return _CheckBlance(_root, blackNum, count);
    }
 
    void _InOrder(Node* root)
    {
        if (root == nullptr)
        {
            return;
        }
 
        _InOrder(root->_left);
        cout << root->_kv.first << ":"<<root->_kv.second<<endl;
        _InOrder(root->_right);
    }
    void InOrder()
    {
        _InOrder(_root);
        cout << endl;
    }
private:
    Node* _root;
};
